#pragma once

#include <common.h>
#include <type.h>
#include <ir/visitor.h>
#include <dot/astdot.h>

namespace ir {
  class Visitor;
}
class AstDot;
namespace ast {
  class NodeBase;

  class NodeCompUnit;
  class NodeGlobalDecl;
  class NodeGlobalDefs;
  class NodeGlobalDef;
  class NodeConstDecl;
  class NodeConstDefs;
  class NodeConstDef;
  class NodeVarDecl;
  class NodeVarDefs;
  class NodeVarDef;
  class NodeArrayConst;
  class NodeArray;
  class NodeArrayInit;
  class NodeFuncDef;
  class NodeFuncDefArgs;
  class NodeBlock;
  class NodeBlockItems;
  class NodeStmt;
  class NodeStmtAssign;
  class NodeStmtReturn;
  class NodeStmtIf;
  class NodeStmtWhile;
  class NodeStmtBreak;
  class NodeStmtContinue;
  class NodeExp;
  class NodeFuncArgs;
  class NodeCond;
}

namespace ast {
  class NodeBase {
  public:
    virtual ~NodeBase();

    virtual void dot_string(std::string & str) = 0;
    virtual void dot_accept(AstDot & ast_dot) = 0;
    virtual void accept(ir::Visitor & visitor) = 0;
  };
  
  class NodeCompUnit : public NodeBase {
  public:
    std::vector<std::variant<
      std::unique_ptr<NodeFuncDef>,
      std::unique_ptr<NodeGlobalDecl>
    >> global_list;

    virtual ~NodeCompUnit() = default;
  
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeGlobalDecl : public NodeBase {
  public:
    VarType type;
    std::unique_ptr<NodeGlobalDefs> global_defs;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeGlobalDefs : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeGlobalDef>> global_def_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeGlobalDef : public NodeBase {
  public:
    std::string iden;
    int int_number;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeConstDecl : public NodeBase {
  public:
    VarType type;
    std::unique_ptr<NodeConstDefs> const_defs;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeConstDefs : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeConstDef>> const_def_list;
  
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeConstDef : public NodeBase {
  public:
    std::string iden;
    int int_number;
  
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeVarDecl : public NodeBase {
  public:
    VarType type;
    std::unique_ptr<NodeVarDefs> var_defs;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeVarDefs : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeVarDef>> var_def_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeVarDef : public NodeBase {
  public:
    enum class Stage {
      EXP_INIT,
      EXP_NOT_INIT,
      ARR_ZERO_INIT,
      ARR_INIT,
      ARR_NOT_INIT,
    } stage;
    std::string iden;
    std::unique_ptr<NodeExp> exp;
    std::unique_ptr<NodeArrayConst> array_const;
    std::unique_ptr<NodeArrayInit> array_init;

    NodeVarDef(Stage stage);
  
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeArrayConst : public NodeBase {
  public:
    std::vector<int> int_number_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeArray : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeExp>> exp_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeArrayInit : public NodeBase {
  public:
    std::vector<std::variant<
      std::unique_ptr<NodeExp>,
      std::unique_ptr<NodeArrayInit>
    >> init_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };
  
  class NodeFuncDef : public NodeBase {
  public:
    enum Stage { HAVE_ARGS, NOT_ARGS } stage;
    VarType type;
    std::string iden;
    std::unique_ptr<NodeFuncDefArgs> func_def_args;
    std::unique_ptr<NodeBlock> block; 

    NodeFuncDef(Stage stage);

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeFuncDefArgs : public NodeBase {
  public:
    // size must is same and not zero
    std::vector<std::string> iden_list;
    std::vector<VarType> type_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeFuncCall : public NodeBase {
  public:
    std::string iden;
    std::optional<std::unique_ptr<NodeFuncArgs>> func_args;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeBlock : public NodeBase {
  public:
    std::unique_ptr<NodeBlockItems> block_item_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeBlockItems : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeStmt>> stmt_list;
    std::vector<std::unique_ptr<NodeVarDecl>> var_decl_list;
    std::vector<std::unique_ptr<NodeConstDecl>> const_decl_list;
  
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };
 
  class NodeStmt : public NodeBase {
  public:
    enum class Stage {
      ASSIGN,
      RETURN,
      IF,
      WHILE,
      BREAK,
      CONTINUE,
      FUNCCALL,
      BLOCK,
    } stage;
    std::unique_ptr<NodeStmtAssign> stmt_assign;
    std::unique_ptr<NodeStmtReturn> stmt_return;
    std::unique_ptr<NodeStmtIf> stmt_if;
    std::unique_ptr<NodeStmtWhile> stmt_while;
    std::unique_ptr<NodeStmtBreak> stmt_break;
    std::unique_ptr<NodeStmtContinue> stmt_continue;
    std::unique_ptr<NodeStmtFuncCall> stmt_func_call;
    std::unique_ptr<NodeBlock> block;

    NodeStmt(Stage stage);

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtAssign : public NodeBase {
  public:
    enum Stage {
      NOARR,
      ARRAY,
    } stage;
    std::string iden;
    std::unique_ptr<NodeArray> array;
    std::unique_ptr<NodeExp> exp;

    NodeStmtAssign(Stage stage);

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtReturn : public NodeBase {
  public:
    std::optional<std::unique_ptr<NodeExp>> exp;
 
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtIf : public NodeBase {
  public:
    enum class Stage {
      NO_ELSE,
      ELSE,
    } stage;
    std::unique_ptr<NodeCond> cond;
    std::unique_ptr<NodeStmt> stmt_true;
    std::unique_ptr<NodeStmt> stmt_false;

    NodeStmtIf(Stage stage);
 
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtWhile : public NodeBase {
  public:
    std::unique_ptr<NodeCond> cond;
    std::unique_ptr<NodeStmt> stmt;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtBreak : public NodeBase {
  public:
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtContinue : public NodeBase {
  public:
    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeStmtFuncCall : public NodeBase {
  public:
    std::unique_ptr<ast::NodeFuncCall> func_call;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  /**
   * Stage      valid
   * EXP        exp1 op_type exp2
   * BRAC_EXP   exp1
   * INT_NUMBER int_number
   * ARRAY      iden array
   * FUNC       iden func_args
   * */
  class NodeExp : public NodeBase {
  public:
    enum Stage {
      EXP,
      BRAC_EXP,
      INT_NUMBER,
      FLOAT_NUMBER,
      IDEN,
      ARRAY,
      FUNC,
    } stage;
    std::unique_ptr<NodeExp> exp1;
    OpType op_type;
    std::unique_ptr<NodeExp> exp2;
    int int_number;
    float float_number;
    std::string iden;
    std::unique_ptr<NodeArray> array;
    std::unique_ptr<NodeFuncCall> func_call;

    NodeExp(Stage stage);

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeFuncArgs : public NodeBase {
  public:
    std::vector<std::unique_ptr<NodeExp>> exp_list;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };

  class NodeCond : public NodeBase {
  public:
    std::unique_ptr<NodeExp> exp;

    virtual void dot_string(std::string & str) override;
    virtual void dot_accept(AstDot & ast_dot) override;
    virtual void accept(ir::Visitor & visitor) override;
  };
}
